# -*- coding: utf-8 -*-
import signal
import sys
import matplotlib.pyplot as plt
import numpy as np
sys.path.append("./pyDmps")
sys.path.append(".")
import dmp_discrete
import LIP
import LIP_motion
import load_trajectory as ld
reload(LIP)
reload(dmp_discrete)
reload(LIP_motion)


def cb(s, f):
    print 'recv signal', s
    sys.exit(0)

# global var
array_e = np.zeros(10000)
array_x = np.zeros(10000)
array_dx = np.zeros(10000)
array_s = np.zeros(10000)
array_state = np.zeros(10000)
array_external_force = np.zeros(2000)
array_swing_x = np.zeros(10000)
array_swing_z = np.zeros(10000)
array_l2 = np.zeros(10000)
array_t2 = np.zeros(10000)
l1 = 0
t1 = 0
i = 0
j = 0
tau_value = 0
standard_tau = 0
step_length = 0.0
data_point_num = 300

alg_flag = input('algorithm flag:\n \
                 1:algorithm in process\n \
                 2:algorithm in begin\n \
                 0:fix gait 35cm\n \
                 -1:fix gait 30cm\n')
if alg_flag == 1:
    print('alg_process is chosen')
    recorder = open("./data/data_process", "w+")

if alg_flag == 2:
    print('alg_begin is chosen')
    recorder = open("./data/data_begin", "w+")

if alg_flag == 0:
    print('alg_no is chosen')
    recorder = open("./data/data_no_fix35", "w+")

if alg_flag == -1:
    print('alg_no is chosen')
    recorder = open("./data/data_no_fix30", "w+")

force_time_begin = 1000
force_time_end = 4000
posed_force = 20.0
simulation_times = 6000


def fun(lip):
    global tau_value, dmp, recorder, i, j, \
        step_length, data_point_num, t1, l1, standard_tau, alg_flag
    i = i+1
    y, dy, ddy = dmp.step(tau=tau_value, external_force=external_force)
    lip.set_swing_foot_pos(y[0], y[1])

    # external force
    force_tmp = 0
    if i < force_time_begin:
        force_tmp = 0
    if i >= force_time_begin and i < force_time_end:
        force_tmp = posed_force
    if i > force_time_end:
        force_tmp = 0
    lip.set_external_force(force_tmp)

    observation_, reward, done, info = lip.update()
    array_e[i] = lip.E
    array_swing_x[i] = lip.stand_leg_world_x + lip.left_foot_x
    array_swing_z[i] = lip.left_foot_z
    recorder.write('%f %f %f %f %f %f\n' % (lip.E,
                                            array_swing_x[i], array_swing_z[i],
                                            lip.x, lip.dx,
                                            force_tmp))

    if alg_flag == 1:
        remain_t = dmp.dt*(data_point_num - dmp.step_number)
        l2, t2 = LIP_motion.get_st_with_e(lip.x, lip.dx, 0.7, 0.3, l1,
                                          remain_t, 8.0, 0.001, 20, 0.001)
        new_tau = remain_t/t2 * standard_tau
        if remain_t > 0.15:
            tau_value = new_tau
        dmp.goal[0] = l2
        array_l2[i] = l2
        array_t2[i] = (t2 - remain_t)**2

    if i > simulation_times or done != 0:
        plt.figure()
        plt.plot(array_e)
        plt.figure()
        plt.plot(array_swing_x, array_swing_z)
        plt.plot(array_t2)
        plt.show()
        recorder.close()
        sys.exit(0)

    if y[1] <= 0.0:
        j += 1
        if lip.stand_leg == 'left_leg':
            lip.switch_support_leg('right_leg')
            dmp.y0[0] = lip.left_foot_x
        else:
            lip.switch_support_leg('left_leg')
            dmp.y0[0] = lip.right_foot_x
        l1, t1 = LIP_motion.get_st_with_e(lip.x0, lip.dx0,
                                          0.7, 0.3, 0.3, 0.3,
                                          0.8, 0.1, 1, 0.001)
        lip.update_orbital_energy()
        LIP_motion.eight_state(lip.x, lip.dx, lip.E)
        dmp.goal[0] = l1  # set the goal_x
        step_length = dmp.goal[0]
        data_point_num = t1/dmp.dt
        tau_value = (1/dmp.dt)/data_point_num
        standard_tau = tau_value

        # alg_flag == 0: chose fix gait
        if alg_flag == 0:
            dmp.goal[0] = 0.35  # set the goal_x
            tau_value = (1/dmp.dt)/250.0

        if alg_flag == -1:
            dmp.goal[0] = 0.3  # set the goal_x
            tau_value = (1/dmp.dt)/250.0
        dmp.reset_state()


if __name__ == "__main__":
    global tau_value
    external_force = np.array([0, 0])
    signal.signal(signal.SIGTERM, cb)
    signal.signal(signal.SIGINT, cb)
    trajectory = ld.loadTrajectory("swing_trajectory.txt")
    trajectory_x = trajectory[:, 0]
    trajectory_y = trajectory[:, 1]
    plt.plot(trajectory_x, trajectory_y, 'r--', lw=2)
    dmp = dmp_discrete.DMPs_discrete(dmps=2, bfs=100)
    dmp.imitate_path(y_des=np.array([trajectory_x, trajectory_y]))
    dmp.goal[0] = 0.2
    dmp.y0[0] = 0.0
    step_length = dmp.goal[0]
    dmp.reset_state()
    data_point_num = 300
    tau_value = (1 / dmp.dt)/data_point_num
    standard_tau = (1 / dmp.dt)/data_point_num
    lip = LIP.LIP(0.1, 0.1, 0.7, 0.001, 'left_leg')
    lip.inster_function = fun
    lip.run()
